#!/usr/bin/env python3

import argparse
from glob import glob
import matplotlib
matplotlib.use('agg')
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.colors import ListedColormap, BoundaryNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import os
import sys
import warnings
warnings.filterwarnings('ignore')

parser = argparse.ArgumentParser(description='Plot 1d advection test cases.', formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(      dest='file_pattern', help='NetCDF esult file pattern')
parser.add_argument('-s', dest='time_step_stride', help='Time step stride', default=5, type=int)
parser.add_argument('-e', dest='end_time_step', help='End time step', type=int)
args = parser.parse_args()

file_paths = sorted(glob(args.file_pattern))

if len(file_paths) == 0:
	print(f'[Error]: Invalid file pattern {args.file_pattern}!')

pdf_file_name = os.path.basename(args.file_pattern.split('.')[0])

colormap = [
	[1.00,0.91,0.91],
	[0.84,0.84,1.00],
	[0.38,0.40,0.98],
	[0.04,0.15,0.89],
	[0.06,0.40,0.46],
	[0.15,0.71,0.14],
	[0.54,0.85,0.17],
	[0.98,0.98,0.22],
	[0.99,0.53,0.14],
	[0.99,0.09,0.11]
]

with PdfPages(pdf_file_name + '.pdf') as pdf:
	for k in range(0, len(file_paths), args.time_step_stride):
		fig = plt.figure(figsize=(8, 8))
		ax = fig.add_subplot(1, 1, 1)
		print(f'==> {file_paths[k]}')
		f = Dataset(file_paths[k], 'r')
		x = f.variables['x'][:]
		y = f.variables['y'][:]
		rho = f.variables['rho'][0,:,:]
		im = ax.contourf(x, y, rho, cmap='rainbow', vmin=0, vmax=1)
		ax.set_xlabel('X')
		ax.set_ylabel('Y')
		ax.set_title(f'{f.scheme.upper()} min: {rho.min()} max: {rho.max()}')
		plt.colorbar(im, make_axes_locatable(ax).append_axes('right', size='5%', pad=0.05))
		f.close()
		pdf.savefig()
		plt.close()

print(f'==> {pdf_file_name}.pdf saved!')